{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NeuromatchAcademy/course-content/blob/master/tutorials/W3D3_NetworkCausality/W3D3_Tutorial3.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "# Neuromatch Academy 2020 -- Week 3 Day 5 Tutorial 3\n",
    "# Causality Day - Simultaneous fitting/regression\n",
    "\n",
    "**Content creators**: Ari Benjamin, Tony Liu, Konrad Kording\n",
    "\n",
    "**Content reviewers**: Mike X Cohen, Madineh Sarvestani, Ella Batty, Michael Waskom"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Tutorial objectives\n",
    "\n",
    "This is tutorial 3 on our day of examining causality. Below is the high level outline of what we'll cover today, with the sections we will focus on in this notebook in bold:\n",
    "\n",
    "1.   Master definitions of causality\n",
    "2.   Understand that estimating causality is possible\n",
    "3.   Learn 4 different methods and understand when they fail\n",
    "    1. perturbations\n",
    "    2. correlations\n",
    "    3. **simultaneous fitting/regression**\n",
    "    4. instrumental variables\n",
    "\n",
    "### Notebook 3 objectives\n",
    "\n",
    "In tutorial 2 we explored correlation as an approximation for causation and learned that correlation $\\neq$ causation for larger networks. However, computing correlations is a rather simple approach, and you may be wondering: will more sophisticated techniques allow us to better estimate causality? Can't we control for things? \n",
    "\n",
    "Here we'll use some common advanced (but controversial) methods that estimate causality from observational data. These methods rely on fitting a function to our data directly, instead of trying to use perturbations or correlations. Since we have the full closed-form equation of our system, we can try these methods and see how well they work in estimating causal connectivity when there are no perturbations. Specifically, we will:\n",
    "\n",
    "- Learn about more advanced (but also controversial) techniques for estimating causality\n",
    "    - conditional probabilities (**regression**)\n",
    "- Explore limitations and failure modes\n",
    "    - understand the problem of **omitted variable bias**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:41.111655Z",
     "iopub.status.busy": "2021-05-25T01:24:41.111046Z",
     "iopub.status.idle": "2021-05-25T01:24:41.836596Z",
     "shell.execute_reply": "2021-05-25T01:24:41.835486Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from sklearn.multioutput import MultiOutputRegressor\n",
    "from sklearn.linear_model import Lasso"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:41.841801Z",
     "iopub.status.busy": "2021-05-25T01:24:41.840775Z",
     "iopub.status.idle": "2021-05-25T01:24:41.925719Z",
     "shell.execute_reply": "2021-05-25T01:24:41.924670Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Figure settings\n",
    "import ipywidgets as widgets       # interactive display\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:41.952959Z",
     "iopub.status.busy": "2021-05-25T01:24:41.933733Z",
     "iopub.status.idle": "2021-05-25T01:24:41.956116Z",
     "shell.execute_reply": "2021-05-25T01:24:41.955609Z"
    }
   },
   "outputs": [],
   "source": [
    "# @title Helper functions\n",
    "\n",
    "\n",
    "def sigmoid(x):\n",
    "    \"\"\"\n",
    "    Compute sigmoid nonlinearity element-wise on x.\n",
    "\n",
    "    Args:\n",
    "        x (np.ndarray): the numpy data array we want to transform\n",
    "    Returns\n",
    "        (np.ndarray): x with sigmoid nonlinearity applied\n",
    "    \"\"\"\n",
    "    return 1 / (1 + np.exp(-x))\n",
    "\n",
    "\n",
    "def logit(x):\n",
    "    \"\"\"\n",
    "\n",
    "    Applies the logit (inverse sigmoid) transformation\n",
    "\n",
    "    Args:\n",
    "        x (np.ndarray): the numpy data array we want to transform\n",
    "    Returns\n",
    "        (np.ndarray): x with logit nonlinearity applied\n",
    "    \"\"\"\n",
    "    return np.log(x/(1-x))\n",
    "\n",
    "\n",
    "def create_connectivity(n_neurons, random_state=42, p=0.9):\n",
    "    \"\"\"\n",
    "    Generate our nxn causal connectivity matrix.\n",
    "\n",
    "    Args:\n",
    "        n_neurons (int): the number of neurons in our system.\n",
    "        random_state (int): random seed for reproducibility\n",
    "\n",
    "    Returns:\n",
    "        A (np.ndarray): our 0.1 sparse connectivity matrix\n",
    "    \"\"\"\n",
    "    np.random.seed(random_state)\n",
    "    A_0 = np.random.choice([0, 1], size=(n_neurons, n_neurons), p=[p, 1 - p])\n",
    "\n",
    "    # set the timescale of the dynamical system to about 100 steps\n",
    "    _, s_vals, _ = np.linalg.svd(A_0)\n",
    "    A = A_0 / (1.01 * s_vals[0])\n",
    "\n",
    "    # _, s_val_test, _ = np.linalg.svd(A)\n",
    "    # assert s_val_test[0] < 1, \"largest singular value >= 1\"\n",
    "\n",
    "    return A\n",
    "\n",
    "\n",
    "def get_regression_estimate_full_connectivity(X):\n",
    "    \"\"\"\n",
    "    Estimates the connectivity matrix using lasso regression.\n",
    "\n",
    "    Args:\n",
    "        X (np.ndarray): our simulated system of shape (n_neurons, timesteps)\n",
    "        neuron_idx (int): optionally provide a neuron idx to compute connectivity for\n",
    "    Returns:\n",
    "        V (np.ndarray): estimated connectivity matrix of shape (n_neurons, n_neurons).\n",
    "                        if neuron_idx is specified, V is of shape (n_neurons,).\n",
    "    \"\"\"\n",
    "    n_neurons = X.shape[0]\n",
    "\n",
    "    # Extract Y and W as defined above\n",
    "    W = X[:, :-1].transpose()\n",
    "    Y = X[:, 1:].transpose()\n",
    "\n",
    "    # apply inverse sigmoid transformation\n",
    "    Y = logit(Y)\n",
    "\n",
    "    # fit multioutput regression\n",
    "    reg = MultiOutputRegressor(Lasso(fit_intercept=False,\n",
    "                                     alpha=0.01, max_iter=250 ), n_jobs=-1)\n",
    "    reg.fit(W, Y)\n",
    "\n",
    "    V = np.zeros((n_neurons, n_neurons))\n",
    "    for i, estimator in enumerate(reg.estimators_):\n",
    "        V[i, :] = estimator.coef_\n",
    "\n",
    "    return V\n",
    "\n",
    "\n",
    "def get_regression_corr_full_connectivity(n_neurons, A, X, observed_ratio, regression_args):\n",
    "    \"\"\"\n",
    "    A wrapper function for our correlation calculations between A and the V estimated\n",
    "    from regression.\n",
    "\n",
    "    Args:\n",
    "        n_neurons (int): number of neurons\n",
    "        A (np.ndarray): connectivity matrix\n",
    "        X (np.ndarray): dynamical system\n",
    "        observed_ratio (float): the proportion of n_neurons observed, must be betweem 0 and 1.\n",
    "        regression_args (dict): dictionary of lasso regression arguments and hyperparameters\n",
    "\n",
    "    Returns:\n",
    "        A single float correlation value representing the similarity between A and R\n",
    "    \"\"\"\n",
    "    assert (observed_ratio > 0) and (observed_ratio <= 1)\n",
    "\n",
    "    sel_idx = np.clip(int(n_neurons*observed_ratio), 1, n_neurons)\n",
    "\n",
    "    sel_X = X[:sel_idx, :]\n",
    "    sel_A = A[:sel_idx, :sel_idx]\n",
    "\n",
    "    sel_V = get_regression_estimate_full_connectivity(sel_X)\n",
    "    return np.corrcoef(sel_A.flatten(), sel_V.flatten())[1,0], sel_V\n",
    "\n",
    "\n",
    "def see_neurons(A, ax, ratio_observed=1, arrows=True):\n",
    "    \"\"\"\n",
    "    Visualizes the connectivity matrix.\n",
    "\n",
    "    Args:\n",
    "        A (np.ndarray): the connectivity matrix of shape (n_neurons, n_neurons)\n",
    "        ax (plt.axis): the matplotlib axis to display on\n",
    "\n",
    "    Returns:\n",
    "        Nothing, but visualizes A.\n",
    "    \"\"\"\n",
    "    n = len(A)\n",
    "\n",
    "    ax.set_aspect('equal')\n",
    "    thetas = np.linspace(0, np.pi * 2, n, endpoint=False)\n",
    "    x, y = np.cos(thetas), np.sin(thetas),\n",
    "    if arrows:\n",
    "      for i in range(n):\n",
    "          for j in range(n):\n",
    "              if A[i, j] > 0:\n",
    "                  ax.arrow(x[i], y[i], x[j] - x[i], y[j] - y[i], color='k', head_width=.05,\n",
    "                          width = A[i, j] / 25,shape='right', length_includes_head=True,\n",
    "                          alpha = .2)\n",
    "    if ratio_observed < 1:\n",
    "      nn = int(n * ratio_observed)\n",
    "      ax.scatter(x[:nn], y[:nn], c='r', s=150, label='Observed')\n",
    "      ax.scatter(x[nn:], y[nn:], c='b', s=150, label='Unobserved')\n",
    "      ax.legend(fontsize=15)\n",
    "    else:\n",
    "      ax.scatter(x, y, c='k', s=150)\n",
    "    ax.axis('off')\n",
    "\n",
    "\n",
    "def simulate_neurons(A, timesteps, random_state=42):\n",
    "    \"\"\"\n",
    "    Simulates a dynamical system for the specified number of neurons and timesteps.\n",
    "\n",
    "    Args:\n",
    "        A (np.array): the connectivity matrix\n",
    "        timesteps (int): the number of timesteps to simulate our system.\n",
    "        random_state (int): random seed for reproducibility\n",
    "\n",
    "    Returns:\n",
    "        - X has shape (n_neurons, timeteps).\n",
    "    \"\"\"\n",
    "    np.random.seed(random_state)\n",
    "\n",
    "\n",
    "    n_neurons = len(A)\n",
    "    X = np.zeros((n_neurons, timesteps))\n",
    "\n",
    "    for t in range(timesteps - 1):\n",
    "        # solution\n",
    "        epsilon = np.random.multivariate_normal(np.zeros(n_neurons), np.eye(n_neurons))\n",
    "        X[:, t + 1] = sigmoid(A.dot(X[:, t]) + epsilon)\n",
    "\n",
    "        assert epsilon.shape == (n_neurons,)\n",
    "    return X\n",
    "\n",
    "\n",
    "def correlation_for_all_neurons(X):\n",
    "  \"\"\"Computes the connectivity matrix for the all neurons using correlations\n",
    "\n",
    "    Args:\n",
    "        X: the matrix of activities\n",
    "\n",
    "    Returns:\n",
    "        estimated_connectivity (np.ndarray): estimated connectivity for the selected neuron, of shape (n_neurons,)\n",
    "        \"\"\"\n",
    "  n_neurons = len(X)\n",
    "  S = np.concatenate([X[:, 1:], X[:, :-1]], axis=0)\n",
    "  R = np.corrcoef(S)[:n_neurons, n_neurons:]\n",
    "  return R\n",
    "\n",
    "\n",
    "def get_sys_corr(n_neurons, timesteps, random_state=42, neuron_idx=None):\n",
    "    \"\"\"\n",
    "    A wrapper function for our correlation calculations between A and R.\n",
    "\n",
    "    Args:\n",
    "        n_neurons (int): the number of neurons in our system.\n",
    "        timesteps (int): the number of timesteps to simulate our system.\n",
    "        random_state (int): seed for reproducibility\n",
    "        neuron_idx (int): optionally provide a neuron idx to slice out\n",
    "\n",
    "    Returns:\n",
    "        A single float correlation value representing the similarity between A and R\n",
    "    \"\"\"\n",
    "\n",
    "    A = create_connectivity(n_neurons, random_state)\n",
    "    X = simulate_neurons(A, timesteps)\n",
    "\n",
    "    R = correlation_for_all_neurons(X)\n",
    "\n",
    "    return np.corrcoef(A.flatten(), R.flatten())[0, 1]\n",
    "\n",
    "\n",
    "def get_regression_corr(n_neurons, A, X, observed_ratio, regression_args, neuron_idx=None):\n",
    "    \"\"\"\n",
    "\n",
    "    A wrapper function for our correlation calculations between A and the V estimated\n",
    "    from regression.\n",
    "\n",
    "    Args:\n",
    "        n_neurons (int): the number of neurons in our system.\n",
    "        A (np.array): the true connectivity\n",
    "        X (np.array): the simulated system\n",
    "        observed_ratio (float): the proportion of n_neurons observed, must be between 0 and 1.\n",
    "        regression_args (dict): dictionary of lasso regression arguments and hyperparameters\n",
    "        neuron_idx (int): optionally provide a neuron idx to compute connectivity for\n",
    "\n",
    "    Returns:\n",
    "        A single float correlation value representing the similarity between A and R\n",
    "    \"\"\"\n",
    "    assert (observed_ratio > 0) and (observed_ratio <= 1)\n",
    "\n",
    "    sel_idx = np.clip(int(n_neurons * observed_ratio), 1, n_neurons)\n",
    "    selected_X = X[:sel_idx, :]\n",
    "    selected_connectivity = A[:sel_idx, :sel_idx]\n",
    "\n",
    "    estimated_selected_connectivity = get_regression_estimate(selected_X, neuron_idx=neuron_idx)\n",
    "    if neuron_idx is None:\n",
    "        return np.corrcoef(selected_connectivity.flatten(),\n",
    "                           estimated_selected_connectivity.flatten())[1, 0], estimated_selected_connectivity\n",
    "    else:\n",
    "        return np.corrcoef(selected_connectivity[neuron_idx, :],\n",
    "                           estimated_selected_connectivity)[1, 0], estimated_selected_connectivity\n",
    "\n",
    "\n",
    "def plot_connectivity_matrix(A, ax=None):\n",
    "  \"\"\"Plot the (weighted) connectivity matrix A as a heatmap\n",
    "\n",
    "    Args:\n",
    "      A (ndarray): connectivity matrix (n_neurons by n_neurons)\n",
    "      ax: axis on which to display connectivity matrix\n",
    "  \"\"\"\n",
    "  if ax is None:\n",
    "    ax = plt.gca()\n",
    "  lim = np.abs(A).max()\n",
    "  ax.imshow(A, vmin=-lim, vmax=lim, cmap=\"coolwarm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Section 1: Regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:41.963877Z",
     "iopub.status.busy": "2021-05-25T01:24:41.963268Z",
     "iopub.status.idle": "2021-05-25T01:24:42.014325Z",
     "shell.execute_reply": "2021-05-25T01:24:42.013729Z"
    },
    "outputId": "36cdf639-7c1d-42df-c8f7-c7feddb97edb"
   },
   "outputs": [],
   "source": [
    "#@title Video 1: Regression approach\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"Av4LaXZdgDo\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "You may be familiar with the idea that correlation only implies causation when there no hidden *confounders*. This aligns with our intuition that correlation only implies causality when no alternative variables could explain away a correlation.\n",
    "\n",
    "**A confounding example**:\n",
    "Suppose you observe that people who sleep more do better in school. It's a nice correlation. But what else could explain it? Maybe people who sleep more are richer, don't work a second job, and have time to actually do homework. If you want to ask if sleep *causes* better grades, and want to answer that with correlations, you have to control for all possible confounds.\n",
    "\n",
    "A confound is any variable that affects both the outcome and your original covariate. In our example, confounds are things that affect both sleep and grades. \n",
    "\n",
    "**Controlling for a confound**: \n",
    "Confonds can be controlled for by adding them as covariates in a regression. But for your coefficients to be causal effects, you need three things:\n",
    " \n",
    "1.   **All** confounds are included as covariates\n",
    "2.   Your regression assumes the same mathematical form of how covariates relate to outcomes (linear, GLM, etc.)\n",
    "3.   No covariates are caused *by* both the treatment (original variable) and the outcome. These are [colliders](https://en.wikipedia.org/wiki/Collider_(statistics)); we won't introduce it today (but Google it on your own time! Colliders are very counterintuitive.)\n",
    "\n",
    "In the real world it is very hard to guarantee these conditions are met. In the brain it's even harder (as we can't measure all neurons). Luckily today we simulated the system ourselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:42.019843Z",
     "iopub.status.busy": "2021-05-25T01:24:42.019151Z",
     "iopub.status.idle": "2021-05-25T01:24:42.055257Z",
     "shell.execute_reply": "2021-05-25T01:24:42.054754Z"
    },
    "outputId": "03512b4c-7c8a-4866-a3d5-cfc9f1aa1011"
   },
   "outputs": [],
   "source": [
    "#@title Video 2: Fitting a GLM\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"GvMj9hRv5Ak\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Section 1.1: Recovering connectivity by model fitting\n",
    "\n",
    "Recall that in our system each neuron effects every other via:\n",
    "\n",
    "$$\n",
    "\\vec{x}_{t+1} = \\sigma(A\\vec{x}_t + \\epsilon_t), \n",
    "$$\n",
    "\n",
    "where $\\sigma$ is our sigmoid nonlinearity from before: $\\sigma(x) = \\frac{1}{1 + e^{-x}}$\n",
    "\n",
    "Our system is a closed system, too, so there are no omitted variables. The regression coefficients should be the causal effect. Are they?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "We will use a regression approach to estimate the causal influence of all neurons to neuron #1. Specifically, we will use linear regression to determine the $A$ in:\n",
    "\n",
    "$$\n",
    "\\sigma^{-1}(\\vec{x}_{t+1}) = A\\vec{x}_t + \\epsilon_t ,\n",
    "$$\n",
    "\n",
    "where $\\sigma^{-1}$ is the inverse sigmoid transformation, also sometimes referred to as the **logit** transformation: $\\sigma^{-1}(x) = \\log(\\frac{x}{1-x})$.\n",
    "\n",
    "Let $W$ be the $\\vec{x}_t$ values, up to the second-to-last timestep $T-1$:\n",
    "\n",
    "$$\n",
    "W = \n",
    "\\begin{bmatrix}\n",
    "\\mid & \\mid & ... & \\mid \\\\ \n",
    "\\vec{x}_0  & \\vec{x}_1  & ... & \\vec{x}_{T-1}  \\\\ \n",
    "\\mid & \\mid & ... & \\mid\n",
    "\\end{bmatrix}_{n \\times (T-1)}\n",
    "$$\n",
    "\n",
    "Let $Y$ be the $\\vec{x}_{t+1}$ values for a selected neuron, indexed by $i$, starting from the second timestep up to the last timestep $T$:\n",
    "\n",
    "$$\n",
    "Y = \n",
    "\\begin{bmatrix}\n",
    "x_{i,1}  & x_{i,2}  & ... & x_{i, T}  \\\\ \n",
    "\\end{bmatrix}_{1 \\times (T-1)}\n",
    "$$\n",
    "\n",
    "You will then fit the following model:\n",
    "\n",
    "$$\n",
    "\\sigma^{-1}(Y^T) = W^TV\n",
    "$$\n",
    "\n",
    "where $V$ is the $n \\times 1$ coefficient matrix of this regression, which will be the estimated connectivity matrix between the selected neuron and the rest of the neurons.\n",
    "\n",
    "**Review**: As you learned Friday of Week 1, *lasso* a.k.a. **$L_1$ regularization** causes the coefficients to be sparse, containing mostly zeros. Think about why we want this here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Exercise 1: Use linear regression plus lasso to estimate causal connectivities\n",
    "\n",
    "You will now create a function to fit the above regression model and V. We will then call this function to examine how close the regression vs the correlation is to true causality.\n",
    "\n",
    "**Code**:\n",
    "\n",
    "You'll notice that we've transposed both $Y$ and $W$ here and in the code we've already provided below. Why is that? \n",
    "\n",
    "This is because the machine learning models provided in scikit-learn expect the *rows* of the input data to be the observations, while the *columns* are the variables. We have that inverted in our definitions of $Y$ and $W$, with the timesteps of our system (the observations) as the columns. So we transpose both matrices to make the matrix orientation correct for scikit-learn.\n",
    "\n",
    "\n",
    "- Because of the abstraction provided by scikit-learn, fitting this regression will just be a call to initialize the `Lasso()` estimator and a call to the `fit()` function\n",
    "- Use the following hyperparameters for the `Lasso` estimator:\n",
    "    - `alpha = 0.01`\n",
    "    - `fit_intercept = False`\n",
    "- How do we obtain $V$ from the fitted model?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:42.062055Z",
     "iopub.status.busy": "2021-05-25T01:24:42.061489Z",
     "iopub.status.idle": "2021-05-25T01:24:44.986785Z",
     "shell.execute_reply": "2021-05-25T01:24:44.987291Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_regression_estimate(X, neuron_idx):\n",
    "    \"\"\"\n",
    "    Estimates the connectivity matrix using lasso regression.\n",
    "\n",
    "    Args:\n",
    "        X (np.ndarray): our simulated system of shape (n_neurons, timesteps)\n",
    "        neuron_idx (int):  a neuron index to compute connectivity for\n",
    "\n",
    "    Returns:\n",
    "        V (np.ndarray): estimated connectivity matrix of shape (n_neurons, n_neurons).\n",
    "                        if neuron_idx is specified, V is of shape (n_neurons,).\n",
    "    \"\"\"\n",
    "    # Extract Y and W as defined above\n",
    "    W = X[:, :-1].transpose()\n",
    "    Y = X[[neuron_idx], 1:].transpose()\n",
    "\n",
    "    # Apply inverse sigmoid transformation\n",
    "    Y = logit(Y)\n",
    "\n",
    "    ############################################################################\n",
    "    ## TODO: Insert your code here to fit a regressor with Lasso. Lasso captures\n",
    "    ## our assumption that most connections are precisely 0.\n",
    "    ## Fill in function and remove\n",
    "    raise NotImplementedError(\"Please complete the regression exercise\")\n",
    "    ############################################################################\n",
    "\n",
    "    # Initialize regression model with no intercept and alpha=0.01\n",
    "    regression = ...\n",
    "\n",
    "    # Fit regression to the data\n",
    "    regression.fit(...)\n",
    "\n",
    "    V = regression.coef_\n",
    "\n",
    "    return V\n",
    "\n",
    "# Parameters\n",
    "n_neurons = 50  # the size of our system\n",
    "timesteps = 10000  # the number of timesteps to take\n",
    "random_state = 42\n",
    "neuron_idx = 1\n",
    "\n",
    "A = create_connectivity(n_neurons, random_state)\n",
    "X = simulate_neurons(A, timesteps)\n",
    "\n",
    "\n",
    "# Uncomment below to test your function\n",
    "# V = get_regression_estimate(X, neuron_idx)\n",
    "\n",
    "#print(\"Regression: correlation of estimated connectivity with true connectivity: {:.3f}\".format(np.corrcoef(A[neuron_idx, :], V)[1, 0]))\n",
    "\n",
    "#print(\"Lagged correlation of estimated connectivity with true connectivity: {:.3f}\".format(get_sys_corr(n_neurons, timesteps, random_state, neuron_idx=neuron_idx)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:44.991063Z",
     "iopub.status.busy": "2021-05-25T01:24:44.989933Z",
     "iopub.status.idle": "2021-05-25T01:24:50.956297Z",
     "shell.execute_reply": "2021-05-25T01:24:50.956763Z"
    },
    "outputId": "d35aeb5d-bb15-4b0d-e2fe-516f901daeed"
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "def get_regression_estimate(X, neuron_idx):\n",
    "    \"\"\"\n",
    "    Estimates the connectivity matrix using lasso regression.\n",
    "\n",
    "    Args:\n",
    "        X (np.ndarray): our simulated system of shape (n_neurons, timesteps)\n",
    "        neuron_idx (int):  a neuron index to compute connectivity for\n",
    "\n",
    "    Returns:\n",
    "        V (np.ndarray): estimated connectivity matrix of shape (n_neurons, n_neurons).\n",
    "                        if neuron_idx is specified, V is of shape (n_neurons,).\n",
    "    \"\"\"\n",
    "    # Extract Y and W as defined above\n",
    "    W = X[:, :-1].transpose()\n",
    "    Y = X[[neuron_idx], 1:].transpose()\n",
    "\n",
    "    # Apply inverse sigmoid transformation\n",
    "    Y = logit(Y)\n",
    "\n",
    "    # Initialize regression model with no intercept and alpha=0.01\n",
    "    regression = Lasso(fit_intercept=False, alpha=0.01)\n",
    "\n",
    "    # Fit regression to the data\n",
    "    regression.fit(W, Y)\n",
    "\n",
    "    V = regression.coef_\n",
    "\n",
    "    return V\n",
    "\n",
    "# Parameters\n",
    "n_neurons = 50  # the size of our system\n",
    "timesteps = 10000  # the number of timesteps to take\n",
    "random_state = 42\n",
    "neuron_idx = 1\n",
    "\n",
    "A = create_connectivity(n_neurons, random_state)\n",
    "X = simulate_neurons(A, timesteps)\n",
    "\n",
    "\n",
    "# Uncomment below to test your function\n",
    "V = get_regression_estimate(X, neuron_idx)\n",
    "\n",
    "print(\"Regression: correlation of estimated connectivity with true connectivity: {:.3f}\".format(np.corrcoef(A[neuron_idx, :], V)[1, 0]))\n",
    "\n",
    "print(\"Lagged correlation of estimated connectivity with true connectivity: {:.3f}\".format(get_sys_corr(n_neurons, timesteps, random_state, neuron_idx=neuron_idx)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "You should find that using regression, our estimated connectivity matrix has a correlation of 0.865 with the true connectivity matrix. With correlation, our estimated connectivity matrix has a correlation of 0.703 with the true connectivity matrix.\n",
    "\n",
    "We can see from these numbers that multiple regression is better than simple correlation for estimating connectivity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Section 2: Omitted Variable Bias\n",
    "\n",
    "If we are unable to observe the entire system, **omitted variable bias** becomes a problem. If we don't have access to all the neurons, and so therefore can't control for them, can we still estimate the causal effect accurately?\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Section 2.1: Visualizing subsets of the connectivity matrix\n",
    "\n",
    "We first visualize different subsets of the connectivity matrix when we observe 75% of the neurons vs 25%.\n",
    "\n",
    "Recall the meaning of entries in our connectivity matrix: $A[i,j] = 1$ means a connectivity **from** neuron $i$ **to** neuron $j$ with strength $1$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 770
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:50.959620Z",
     "iopub.status.busy": "2021-05-25T01:24:50.959043Z",
     "iopub.status.idle": "2021-05-25T01:24:52.340671Z",
     "shell.execute_reply": "2021-05-25T01:24:52.341256Z"
    },
    "outputId": "1c5df52e-3ce1-4560-bec3-43d82607580e"
   },
   "outputs": [],
   "source": [
    "#@markdown Execute this cell to visualize subsets of connectivity matrix\n",
    "\n",
    "# Run this cell to visualize the subsets of variables we observe\n",
    "n_neurons = 25\n",
    "A = create_connectivity(n_neurons)\n",
    "\n",
    "fig, axs = plt.subplots(2, 2, figsize=(10, 10))\n",
    "ratio_observed = [0.75, 0.25]  # the proportion of neurons observed in our system\n",
    "\n",
    "for i, ratio in enumerate(ratio_observed):\n",
    "    sel_idx = int(n_neurons * ratio)\n",
    "\n",
    "    offset = np.zeros((n_neurons, n_neurons))\n",
    "    axs[i,1].title.set_text(\"{}% neurons observed\".format(int(ratio * 100)))\n",
    "    offset[:sel_idx, :sel_idx] =  1 + A[:sel_idx, :sel_idx]\n",
    "    im = axs[i, 1].imshow(offset, cmap=\"coolwarm\", vmin=0, vmax=A.max() + 1)\n",
    "    axs[i, 1].set_xlabel(\"Connectivity from\")\n",
    "    axs[i, 1].set_ylabel(\"Connectivity to\")\n",
    "    plt.colorbar(im, ax=axs[i, 1], fraction=0.046, pad=0.04)\n",
    "    see_neurons(A,axs[i, 0],ratio)\n",
    "\n",
    "plt.suptitle(\"Visualizing subsets of the connectivity matrix\", y = 1.05)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Section 2.2: Effects of partial observability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:52.346922Z",
     "iopub.status.busy": "2021-05-25T01:24:52.346358Z",
     "iopub.status.idle": "2021-05-25T01:24:52.400640Z",
     "shell.execute_reply": "2021-05-25T01:24:52.401264Z"
    },
    "outputId": "b3a87e09-75a6-4ca5-833e-115668543851"
   },
   "outputs": [],
   "source": [
    "#@title Video 3: Omitted variable bias\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"5CCib6CTMac\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**Video correction**: the labels \"connectivity from\"/\"connectivity to\" are swapped in the video but fixed in the figures/demos below"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Interactive Demo: Regression performance as a function of the number of observed neurons\n",
    "\n",
    "We will first change the number of observed neurons in the network and inspect the resulting estimates of connectivity in this interactive demo. How does the estimated connectivity differ?\n",
    "\n",
    "**Note:** the plots will take a moment or so to update after moving the slider."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 402,
     "referenced_widgets": [
      "3453bbed7cf944e8911cd45065a2c365",
      "c7df0d2581354dc69bfc6f8262cb689e",
      "e7033aff57e546d09cdc30d86207ce54",
      "5312883a23a84981aee827501f019423",
      "20260b2cd4ad445abe389ae1c937b41a",
      "e68b26f6784e4c9db12b62b0a2eaffc1",
      "681bfbca6624446c8f76f5c537bd9038"
     ]
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:52.416918Z",
     "iopub.status.busy": "2021-05-25T01:24:52.412033Z",
     "iopub.status.idle": "2021-05-25T01:24:55.229197Z",
     "shell.execute_reply": "2021-05-25T01:24:55.228683Z"
    },
    "outputId": "1650216a-c251-4711-a1d9-858c40101e00"
   },
   "outputs": [],
   "source": [
    "#@markdown Execute this cell to enable demo\n",
    "n_neurons = 50\n",
    "A = create_connectivity(n_neurons, random_state=42)\n",
    "X = simulate_neurons(A, 4000, random_state=42)\n",
    "\n",
    "reg_args = {\n",
    "    \"fit_intercept\": False,\n",
    "    \"alpha\": 0.001\n",
    "}\n",
    "\n",
    "@widgets.interact\n",
    "def plot_observed(n_observed=(5, 45, 5)):\n",
    "  to_neuron = 0\n",
    "  fig, axs = plt.subplots(1, 3, figsize=(15, 5))\n",
    "  sel_idx = n_observed\n",
    "  ratio = (n_observed) / n_neurons\n",
    "  offset = np.zeros((n_neurons, n_neurons))\n",
    "  axs[0].title.set_text(\"{}% neurons observed\".format(int(ratio * 100)))\n",
    "  offset[:sel_idx, :sel_idx] =  1 + A[:sel_idx, :sel_idx]\n",
    "  im = axs[1].imshow(offset, cmap=\"coolwarm\", vmin=0, vmax=A.max() + 1)\n",
    "  plt.colorbar(im, ax=axs[1], fraction=0.046, pad=0.04)\n",
    "\n",
    "  see_neurons(A,axs[0], ratio, False)\n",
    "  corr, R =  get_regression_corr_full_connectivity(n_neurons,\n",
    "                                  A,\n",
    "                                  X,\n",
    "                                  ratio,\n",
    "                                  reg_args)\n",
    "\n",
    "  #rect = patches.Rectangle((-.5,to_neuron-.5),n_observed,1,linewidth=2,edgecolor='k',facecolor='none')\n",
    "  #axs[1].add_patch(rect)\n",
    "  big_R = np.zeros(A.shape)\n",
    "  big_R[:sel_idx, :sel_idx] =  1 + R\n",
    "  #big_R[to_neuron, :sel_idx] =  1 + R\n",
    "  im = axs[2].imshow(big_R, cmap=\"coolwarm\", vmin=0, vmax=A.max() + 1)\n",
    "  plt.colorbar(im, ax=axs[2],fraction=0.046, pad=0.04)\n",
    "  c = 'w' if n_observed<(n_neurons-3) else 'k'\n",
    "  axs[2].text(0,n_observed+3,\"Correlation : {:.2f}\".format(corr), color=c, size=15)\n",
    "  #axs[2].axis(\"off\")\n",
    "  axs[1].title.set_text(\"True connectivity\")\n",
    "  axs[1].set_xlabel(\"Connectivity from\")\n",
    "  axs[1].set_ylabel(\"Connectivity to\")\n",
    "\n",
    "  axs[2].title.set_text(\"Estimated connectivity\")\n",
    "  axs[2].set_xlabel(\"Connectivity from\")\n",
    "  #axs[2].set_ylabel(\"Connectivity to\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Next, we will inspect a plot of the correlation between true and estimated connectivity matrices vs the percent of neurons observed over multiple trials.\n",
    "What is the relationship that you see between performance and the number of neurons observed?\n",
    "\n",
    "**Note:** the cell below will take about 25-30 seconds to run.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 483
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:24:55.241197Z",
     "iopub.status.busy": "2021-05-25T01:24:55.240584Z",
     "iopub.status.idle": "2021-05-25T01:25:06.466094Z",
     "shell.execute_reply": "2021-05-25T01:25:06.466566Z"
    },
    "outputId": "d97b251e-95aa-408e-f481-e495d5085a19"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "#@markdown Plot correlation vs. subsampling\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# we'll simulate many systems for various ratios of observed neurons\n",
    "n_neurons = 50\n",
    "timesteps = 5000\n",
    "ratio_observed = [1, 0.75, 0.5, .25, .12]  # the proportion of neurons observed in our system\n",
    "n_trials = 3  # run it this many times to get variability in our results\n",
    "\n",
    "reg_args = {\n",
    "    \"fit_intercept\": False,\n",
    "    \"alpha\": 0.001\n",
    "}\n",
    "\n",
    "corr_data = np.zeros((n_trials, len(ratio_observed)))\n",
    "for trial in range(n_trials):\n",
    "\n",
    "  A = create_connectivity(n_neurons, random_state=trial)\n",
    "  X = simulate_neurons(A, timesteps)\n",
    "  print(\"simulating trial {} of {}\".format(trial + 1, n_trials))\n",
    "\n",
    "\n",
    "  for j, ratio in enumerate(ratio_observed):\n",
    "      result,_ = get_regression_corr_full_connectivity(n_neurons,\n",
    "                                    A,\n",
    "                                    X,\n",
    "                                    ratio,\n",
    "                                    reg_args)\n",
    "      corr_data[trial, j] = result\n",
    "\n",
    "corr_mean = np.nanmean(corr_data, axis=0)\n",
    "corr_std = np.nanstd(corr_data, axis=0)\n",
    "\n",
    "plt.plot(np.asarray(ratio_observed) * 100, corr_mean)\n",
    "plt.fill_between(np.asarray(ratio_observed) * 100,\n",
    "                    corr_mean - corr_std,\n",
    "                    corr_mean + corr_std,\n",
    "                    alpha=.2)\n",
    "plt.xlim([100, 10])\n",
    "plt.xlabel(\"Percent of neurons observed\")\n",
    "plt.ylabel(\"connectivity matrices correlation\")\n",
    "plt.title(\"Performance of regression as a function of the number of neurons observed\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Summary\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:25:06.472299Z",
     "iopub.status.busy": "2021-05-25T01:25:06.471646Z",
     "iopub.status.idle": "2021-05-25T01:25:06.515622Z",
     "shell.execute_reply": "2021-05-25T01:25:06.516064Z"
    },
    "outputId": "0c5f6d74-1635-45d9-d3f2-e243743a7ed3"
   },
   "outputs": [],
   "source": [
    "#@title Video 4: Summary\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"T1uGf1H31wE\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "In this tutorial, we explored:\n",
    "\n",
    "1) Using regression for estimating causality\n",
    "\n",
    "2) The problem of ommitted variable bias, and how it arises in practice"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "W3D5_Tutorial3",
   "provenance": [],
   "toc_visible": true
  },
  "kernel": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
